/**
<code>
  <author name="Zealic" email="zealic(at)gmail.com" type="Creator"/>
</code>
**/
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Net;
using System.Net.NetworkInformation;
using System.Runtime.InteropServices;


namespace Zealib.Win32.Network
{
    /// <summary>
    /// TCP manager
    /// </summary>
    public static class TcpConnectionManager
    {
        #region PInvoke define
        private const int TCP_TABLE_OWNER_PID_ALL = 5;

        [DllImport("iphlpapi.dll", SetLastError = true)]
        private static extern uint GetExtendedTcpTable(
            IntPtr pTcpTable, ref int dwOutBufLen, bool sort, int ipVersion, int tblClass, int reserved);

        [DllImport("iphlpapi.dll")]
        private static extern int SetTcpEntry(ref MIB_TCPROW pTcpRow);


        [StructLayout(LayoutKind.Sequential)]
        private struct MIB_TCPROW
        {
            public TcpState dwState;
            public int dwLocalAddr;
            public int dwLocalPort;
            public int dwRemoteAddr;
            public int dwRemotePort;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct MIB_TCPROW_OWNER_PID
        {
            public TcpState dwState;
            public uint dwLocalAddr;
            public int dwLocalPort;
            public uint dwRemoteAddr;
            public int dwRemotePort;
            public int dwOwningPid;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct MIB_TCPTABLE_OWNER_PID
        {
            public uint dwNumEntries;
            private MIB_TCPROW_OWNER_PID table;
        }
        #endregion

        private static MIB_TCPROW_OWNER_PID[] GetAllTcpConnections()
        {
            const int NO_ERROR = 0;
            const int IP_v4 = 2;
            MIB_TCPROW_OWNER_PID[] tTable = null;
            int buffSize = 0;
            GetExtendedTcpTable(IntPtr.Zero, ref buffSize, true, IP_v4, TCP_TABLE_OWNER_PID_ALL, 0);
            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);
            try
            {
                if (NO_ERROR != GetExtendedTcpTable(buffTable, ref buffSize, true, IP_v4, TCP_TABLE_OWNER_PID_ALL, 0)) return null;
                MIB_TCPTABLE_OWNER_PID tab =
                    (MIB_TCPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_TCPTABLE_OWNER_PID));
                IntPtr rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(tab.dwNumEntries));
                tTable = new MIB_TCPROW_OWNER_PID[tab.dwNumEntries];

                int rowSize = Marshal.SizeOf(typeof(MIB_TCPROW_OWNER_PID));
                for (int i = 0; i < tab.dwNumEntries; i++)
                {
                    MIB_TCPROW_OWNER_PID tcpRow =
                        (MIB_TCPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_TCPROW_OWNER_PID));
                    tTable[i] = tcpRow;
                    rowPtr = (IntPtr)((int)rowPtr + rowSize);
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffTable);
            }
            return tTable;
        }

        private static int TranslatePort(int port)
        {
            return ((port & 0xFF) << 8 | (port & 0xFF00) >> 8);
        }

        public static bool Kill(TcpConnectionInfo conn)
        {
            if (conn == null) throw new ArgumentNullException("conn");
            MIB_TCPROW row = new MIB_TCPROW();
            row.dwState = TcpState.DeleteTcb;
#pragma warning disable 612,618
            row.dwLocalAddr = (int)conn.LocalEndPoint.Address.Address;
            row.dwRemoteAddr = (int)conn.RemoteEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwLocalPort = TranslatePort(conn.LocalEndPoint.Port);
            row.dwRemotePort = TranslatePort(conn.RemoteEndPoint.Port);
            return SetTcpEntry(ref row) == 0;
        }

        public static bool Kill(IPEndPoint localEndPoint, IPEndPoint remoteEndPoint)
        {
            if (localEndPoint == null) throw new ArgumentNullException("localEndPoint");
            if (remoteEndPoint == null) throw new ArgumentNullException("remoteEndPoint");
            MIB_TCPROW row = new MIB_TCPROW();
            row.dwState = TcpState.DeleteTcb;
#pragma warning disable 612,618
            row.dwLocalAddr = (int)localEndPoint.Address.Address;
            row.dwRemoteAddr = (int)remoteEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwLocalPort = TranslatePort(localEndPoint.Port);
            row.dwRemotePort = TranslatePort(remoteEndPoint.Port);
            return SetTcpEntry(ref row) == 0;
        }

        public static TcpConnectionInfo[] GetTable()
        {
            MIB_TCPROW_OWNER_PID[] tcpRows = GetAllTcpConnections();
            if (tcpRows == null) return null;
            List<TcpConnectionInfo> list = new List<TcpConnectionInfo>();
            foreach (MIB_TCPROW_OWNER_PID row in tcpRows)
            {
                int localPort = TranslatePort(row.dwLocalPort);
                int remotePort = TranslatePort(row.dwRemotePort);
                TcpConnectionInfo conn =
                    new TcpConnectionInfo(
                        new IPEndPoint(row.dwLocalAddr, localPort),
                        new IPEndPoint(row.dwRemoteAddr, remotePort),
                        row.dwState);
                list.Add(conn);
            }
            return list.ToArray();
        }

        public static TcpConnectionInfo[] GetTableByProcess(int pid)
        {
            if(pid > 0) throw new ArgumentOutOfRangeException("pid");
            MIB_TCPROW_OWNER_PID[] tcpRows = GetAllTcpConnections();
            if (tcpRows == null) return null;
            List<TcpConnectionInfo> list = new List<TcpConnectionInfo>();
            foreach (MIB_TCPROW_OWNER_PID row in tcpRows)
            {
                if (row.dwOwningPid == pid)
                {
                    int localPort = TranslatePort(row.dwLocalPort);
                    int remotePort = TranslatePort(row.dwRemotePort);
                    TcpConnectionInfo conn =
                        new TcpConnectionInfo(
                            new IPEndPoint(row.dwLocalAddr, localPort),
                            new IPEndPoint(row.dwRemoteAddr, remotePort),
                            row.dwState);
                    list.Add(conn);
                }
            }
            return list.ToArray();
        }

        public static TcpConnectionInfo[] GetTalbeByCurrentProcess()
        {
            return GetTableByProcess(Process.GetCurrentProcess().Id);
        }

    }
}